import os
from PIL import Image
import torch
import torch.utils.data as data
import torchvision.transforms as transforms
import numpy as np
import matplotlib.pyplot as plt
import albumentations as A
import cv2
'''
albumentations的输入维度是HWC，
transforms.ToTensor()#会将将 HWC 的图像格式（或者numpy）转为 CHW 的 tensor 格式

再输入到torchvision.transforms
'''

class SkinDataset(data.Dataset):
    """
    dataloader for skin lesion segmentation tasks
    """

    def __init__(self, image_root, gt_root):
        # self.images = np.load(image_root)
        # self.gts = np.load(gt_root)
        # self.size = len(self.images)

        self.img_transform = transforms.Compose([
            transforms.ToTensor(),
            #这里需要提前计算每个通道的均值和方差，再在这里设置
            transforms.Normalize([0.485],
                                [0.229])
        ])
        self.gt_transform = transforms.Compose([
            transforms.ToTensor()#会将将 HWC 的图像格式转为 CHW 的 tensor 格式
        ])

        self.transform = A.Compose(#transforms.Compose() 用于整合一系列的图像变换函数，将图片按照 Compose() 中的顺序依次处理。
            [
                A.ShiftScaleRotate(shift_limit = 0.15, scale_limit = 0.15, rotate_limit = 25, p = 0.5, border_mode = 0),#平移缩放旋转
                # A.ColorJitter(),#随机改变图像的亮度、对比度、饱和度和色调
                A.HorizontalFlip(),#用于水平地随机翻转一半图像
                A.VerticalFlip()#用于垂直地随机翻转一半图像
            ]
        )

    def __getitem__(self, index):
        image = np.ones((256,256,1))
        gt = np.ones((256,256,1))

        gt = gt / 255.0

        transformed = self.transform(image = image, mask = gt)#调用albumentations的数据变换
        # transformed['image'] = np.transpose(transformed['image'], (2,0,1))
        # transformed['mask']= np.transpose(transformed['mask'], (2,0,1))

        # transformed['image'] = torch.tensor(transformed['image'])
        # transformed['mask'] = torch.tensor(transformed['mask'])

        image = self.img_transform(transformed['image'])#调用torchvision的数据变换
        gt = self.gt_transform(transformed['mask'])#调用torchvision的数据变换
        return image, gt

    def __len__(self):
        return 8


def get_loader(image_root, gt_root, batchsize, shuffle = True, num_workers = 4, pin_memory = True):
    dataset = SkinDataset(image_root, gt_root)
    data_loader = data.DataLoader(dataset = dataset,
                                  batch_size = batchsize,
                                  shuffle = shuffle,
                                  num_workers = num_workers,
                                  pin_memory = pin_memory)
    return data_loader


class test_dataset:
    def __init__(self, image_root, gt_root):
        # self.images = np.load(image_root)
        # self.gts = np.load(gt_root)

        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize([0.485],
                                 [0.229])
        ])
        self.gt_transform = transforms.ToTensor()

        self.index = 0

    def load_data(self):
        image = np.ones((256, 256, 1))
        gt = np.ones((256, 256, 1))
        gt = gt / 255.0
        self.index += 1

        return image, gt


if __name__ == '__main__':
    dl=get_loader('','',4)
    i=0
    for img,tag in dl:
        print(img.shape,tag.shape,"**")
        i=i+1
    print(i)
    # path = 'data/'
    # tt = SkinDataset(path + 'data_train.npy', path + 'mask_train.npy')
    # for i in range(50):
    #     img, gt = tt.__getitem__(i)
    #
    #     img = torch.transpose(img, 0, 1)
    #     img = torch.transpose(img, 1, 2)
    #     img = img.numpy()
    #     gt = gt.numpy()
    #
    #     plt.imshow(img)
    #     plt.savefig('vis/'+str(i)+".jpg")
    #
    #     plt.imshow(gt[0])
    #     plt.savefig('vis/'+str(i)+'_gt.jpg')
